package com.treeroot.treenet.protocol.abs.tcp;

import com.treeroot.treenet.config.ConfigUtil;
import com.treeroot.treenet.protocol.BaseProtocolHandle;
import com.treeroot.treenet.protocol.abs.tcp.route.DefaultMessageRouter;
import com.treeroot.treenet.protocol.abs.tcp.route.MessageRouter;
import com.treeroot.treenet.protocol.abs.tcp.route.PingHandler;
import com.treeroot.treenet.protocol.abs.tcp.session.TcpSession;
import com.treeroot.treenet.protocol.abs.tcp.session.TcpSessionManager;
import lombok.Getter;
import lombok.NonNull;
import lombok.Setter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.Channel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.util.Iterator;
import java.util.Set;
import java.util.function.Consumer;

/**
 * Description: 抽象 TCP 连接实现类
 * Create By TreeRoot
 * Create Time: 2025/06/02 07:06
 * Class Name: AbsTcpConnection
 *
 * @author TreeRoot
 */
@Getter
public abstract class AbsTcpConnection implements BaseProtocolHandle<TcpSession>, AbstractBaseTcpConnection {

    private static final Logger log = LoggerFactory.getLogger(AbsTcpConnection.class);

    @Setter
    private volatile boolean running = false;
    /**
     * 选择器，用于监听通道事件
     */
    private Selector selector;
    /**
     * 服务器套接字通道，用于监听客户端连接请求
     */
    private ServerSocketChannel serverSocketChannel;

    /**
     * 管理所有连接的会话管理器，超时时间 60秒
     */
    private final TcpSessionManager sessionManager = new TcpSessionManager(60_000);

    @Override
    public void handleRequest(@NonNull Consumer<TcpSession> acceptConnectionConsumer) {
        try {
            selector = Selector.open();
            serverSocketChannel = ServerSocketChannel.open();
            serverSocketChannel.configureBlocking(false);
            serverSocketChannel.bind(new InetSocketAddress(ConfigUtil.getTcpPort()), 1024);
            serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

            running = true;
            log.info("TCP 服务器启动，监听端口：{}", ConfigUtil.getTcpPort());
            sessionManager.start();
            log.info("TCP 会话管理器已启动");
            MessageRouter router = new DefaultMessageRouter();
            router.registerHandler("PING", new PingHandler());
            log.info("[Router] 注册消息类型处理器: PING");
            while (running) {
                //最多等待 1 秒，看看有没有通道发生事件
                int selectCount = selector.select(1000);
                if (selectCount == 0) {
                    continue;
                }
                Set<SelectionKey> selectedKeys = selector.selectedKeys();
                Iterator<SelectionKey> iterator = selectedKeys.iterator();

                while (iterator.hasNext()) {
                    SelectionKey key = iterator.next();
                    // 一定要先移除，防止ConcurrentModificationException
                    iterator.remove();
                    if (!key.isValid()) {
                        continue;
                    }
                    try {
                        if (key.isAcceptable()) {
                            TcpSession tcpSession = new TcpSession(serverSocketChannel, selector, key);
                            tcpSession.setRouter(router);
                            // 连接前钩子
                            acceptBefore(tcpSession);
                            // 执行连接事件处理
                            acceptConnectionConsumer.accept(tcpSession);
                            // 把 session 附加到 key 上
                            // 把新连接加入管理
                            sessionManager.addSession(tcpSession);
                            // 连接后钩子
                            acceptAfter(tcpSession);
                        } else if (key.isReadable()) {
                            //从 客户端读取数据的key中读取key
                            TcpSession tcpSession = (TcpSession) key.attachment();
                            if (tcpSession == null) {
                                log.warn("读事件但未附加 session");
                                continue;
                            }
                            readDataFromClient(key, () -> sessionManager.removeSession(tcpSession.getChannelId()));
                        }
                    } catch (Exception e) {
                        log.error("事件处理异常，关闭通道", e);
                        cancelKeyAndCloseChannel(key);
                        sessionManager.removeSession(sessionManager.findSessionByKey(key).getChannelId());
                    }
                }
            }
        } catch (IOException e) {
            log.error("TCP 服务器异常", e);
        } finally {
            closeServer();
            sessionManager.stop();
        }
    }


    /**
     * 关闭服务器
     */
    private void closeServer() {
        running = false;
        try {
            if (selector != null) {
                for (SelectionKey key : selector.keys()) {
                    cancelKeyAndCloseChannel(key);
                }
                selector.close();
            }
            if (serverSocketChannel != null) {
                serverSocketChannel.close();
            }
            log.info("TCP 服务器已关闭");
        } catch (IOException e) {
            log.error("关闭 TCP 服务器异常", e);
        }
    }

    /**
     * 取消key并关闭通道
     *
     * @param key 钥匙
     */
    private void cancelKeyAndCloseChannel(SelectionKey key) {
        if (key == null) {
            return;
        }
        try {
            key.cancel();
            Channel channel = key.channel();
            if (channel != null) {
                channel.close();
            }
        } catch (IOException e) {
            log.error("关闭通道失败", e);
        }
    }

    /**
     * 接受请求之前
     *
     * @param baseTcpVo 基本tcp vo
     */
    public abstract void acceptBefore(TcpSession baseTcpVo);

    /**
     * 接受请求之后
     *
     * @param baseTcpVo 基本tcp vo
     * @throws IOException IOException
     */
    public abstract void acceptAfter(TcpSession baseTcpVo) throws IOException;

}